MyPath = joinpath(homedir(), "Dropbox/Projects/apr")


using Statistics, LinearAlgebra, StatsBase, Printf, SparseArrays, Plots, SpecialFunctions, InfinitesimalGenerators, EconPDEs, Optim
blue = colorant"rgb(68, 118, 170)"
red = colorant"rgb(204, 102, 119)"
yellow = colorant"rgb(221, 170, 51)"
pgfplotsx()
include("$(MyPath)/code/julia/IncompleteMarkets.jl")

#=============================================================================================

Explore baseline calibration

=============================================================================================#

m = IncompleteMarketsModel(σR = 0.0)

open("$(MyPath)/Tables/TableA7.tex", "w") do myfile
    @printf(myfile, "\\textit{Income process} &  &  \\\\ \n")
    @printf(myfile, "\\quad Labor income growth & \$ \\mu\$ & %.2f  \\\\ \n", m.μ - 0.5 * m.σ^2/2)
    @printf(myfile, "\\quad Labor income volatility & \$ \\sigma\$ & %.2f  \\\\ \n", m.σ)
    @printf(myfile, "\\quad Income multiplier when retired & \$ \\chi_R\$ & %.2f  \\\\ \n", m.χR)
    @printf(myfile, "\\quad Transition probability to unemployment & \$ \\lambda_{EU}\$ & %.2f  \\\\ \n", 1 - exp(-m.λEU))
    @printf(myfile, "\\quad Transition probability to employment & \$ \\lambda_{UE}\$ & %.2f  \\\\ \n", 1 - exp(-m.λUE))
    @printf(myfile, "\\quad Income multiplier when unemployed & \$ \\chi_U\$ & %.2f  \\\\ \n", m.χU)
    @printf(myfile, "\\textit{Asset prices} &  &  \\\\ \n")
    @printf(myfile, "\\quad Interest rate & \$ \\log R \$ & %.2f  \\\\ \n", exp(m.r)-1)
    @printf(myfile, "\\textit{Household preferences} &  &  \\\\ \n")
    @printf(myfile, "\\quad Relative risk aversion & \$ \\gamma\$ & %.2f  \\\\ \n", m.γ)
    @printf(myfile, "\\quad Impatience parameter & \$ \\beta\$ & %.2f  \\\\ \n", exp(-m.ρ))
    @printf(myfile, "\\quad Bequest preference & \$ b\$ & %.2f  \\\\ \n", m.b)
end

# set up state and time grid
println("1. Solving the model")
stategrid = OrderedDict(:w => range(m.wmin^(1/3), m.wmax^(1/3), length = 1000).^3)
τs_old = range(65, 85, step = 0.25)
τs_young = range(20, first(τs_old), step = 0.25)

# specify terminal value
yend = OrderedDict(
    :pR => m.b^(1/(1-m.γ)) .* (m.r + (m.ρ - m.r) / m.γ)^(1/(1 - 1 / m.γ)) * (stategrid[:w] .+ 0.01)
    )

# compute the policy functions in the model using backward iterations
result = solve_model(m, stategrid, τs_young,  τs_old, yend; verbose = false)

# compute fixed point for density
println("2. Computing wealth density as fixed point (asset bequeathed = asset received)")
gE, gU, gR, gYE, gYU, gYR = solve_density(m, stategrid, τs_young, τs_old, result; tol = 1e-4)
g = sum(gE, dims = 2) .+ sum(gU, dims = 2) .+ sum(gR, dims = 2)
g = vec(g ./ sum(g))
G = cumsum(g)
irange = 1:searchsortedfirst(G, 0.99)



# plot asset purchases across different states
purchases = [(m.r .* stategrid[:w] .+ 1 .- vec(mean(result[:cE], dims = 2))) (m.r .* stategrid[:w] .+ m.χU .- vec(mean(result[:cU], dims = 2))) (m.r .* stategrid[:w] .+ m.χR .- vec(mean(result[:cR], dims = 2)))]
p = plot(stategrid[:w][irange], [purchases[irange, 1] purchases[irange, 2] purchases[irange, 3]], linewidth = [2 2 2], linestyle = [:solid :dash :dot],  color = [blue red yellow], lab = ["Employed"  "Unemployed" "Retired"], xlabel = "Financial wealth", ylabel = "Asset purchases", xtickfontsize=12, ytickfontsize=12, guidefontsize = 12, legendfontsize = 12, ylims = (- 1.2, 0.52),  yticks = [-1, -0.75, -0.5, -0.25, 0, 0.25, 0.5], legendfonthalign = :left, legend = :topright)
hline!([0], linestyle=:solid, color = :black, lab = "")
savefig(p, "$(MyPath)/Figures/FigureA7a.pdf") 

# plot density rescaled by population share
p = plot(stategrid[:w][irange], [mean(gE, dims = 2)[irange] ./ sum(mean(gE, dims = 2)) mean(gU, dims = 2)[irange] ./ sum(mean(gU, dims = 2)) mean(gR, dims = 2)[irange] ./ sum(mean(gR, dims = 2))], linewidth = [2 2 2], linestyle = [:solid :dash :dot],  color = [blue red yellow], lab = ["Employed"  "Unemployed" "Retired"], xlabel = "Financial wealth", ylabel = "Density", xtickfontsize=12, ytickfontsize=12, guidefontsize = 12, legendfontsize = 12, legendfonthalign = :left, legend = :topright)
savefig(p, "$(MyPath)/Figures/FigureA7b.pdf") 

# ratio marginal utility
p = plot(stategrid[:w][irange], vec(mean(result[:cU].^(-m.γ) ./ result[:cE].^(-m.γ), dims = 2))[irange], linewidth = 2, linestyle = :solid , color = :black, lab = "", xlabel = "Financial wealth", ylabel = "Ratio marginal utility  U/E", xtickfontsize=12, ytickfontsize=12, guidefontsize = 12, legendfontsize = 12, legendfonthalign = :left, legend = :topright)
savefig(p, "$(MyPath)/Figures/FigureA7c.pdf") 

# average labor income of 50K -> starting salary of 41K
Ybar = compute_laborincome(m,  stategrid, τs_young, τs_old, gYE, gYU, gYR)

# compute welfare gains
println("3. Computing welfare gains")

ϕ = -log(0.95)
dr = - 0.05
dp = - dr / (m.r + ϕ)
ts = 0:50
p = plot(ts, exp.(-ϕ.* ts) .* dr, linewidth = 2, linestyle = :solid, color = :black, ylabel = "Deviation in interest rate",  lab = "",  xlabel = "Years", xtickfontsize=12, ytickfontsize=12, guidefontsize = 12, legendfontsize = 12, legendfonthalign = :left, legend = :topleft)
hline!([0], linestyle=:solid, color = :black, lab = "")
savefig(p, "$(MyPath)/Figures/FigureA8a.pdf") 

wgR_baseline, wgE_baseline, wgU_baseline, wgR_stochastic, wgE_stochastic, wgU_stochastic = solve_welfare(m, stategrid, τs_young, τs_old, result; ϕ = ϕ)
wgEU_baseline = mean(wgE_baseline .* gE .+ wgU_baseline .* gU, dims = 2) ./ mean(gE .+ gU, dims = 2)
wgEU_stochastic = mean(wgE_stochastic .* gE .+ wgU_stochastic .* gU, dims = 2) ./ mean(gE .+ gU, dims = 2)
p = plot(stategrid[:w][irange], [dp .* wgEU_baseline[irange] ./ Ybar  dp .* wgEU_stochastic[irange] ./ Ybar], linewidth = [2 2 2 2 ], linestyle = [:solid :dot],  color = :black, lab = ["Discounting by interest rate" "Discounting by individual MRS" ], xlabel = "Financial wealth", ylabel = "Welfare gain", xtickfontsize=12, ytickfontsize=12, guidefontsize = 12, legendfontsize = 12, legendfonthalign = :left, legend = :topleft)
hline!([0], linestyle=:solid, color = :black, lab = "")
savefig(p, "$(MyPath)/Figures/FigureA8b.pdf") 

wgR_baseline_mean = sum(wgR_baseline .* gYR, dims = 1)
wgR_stochastic_mean = sum(wgR_stochastic .* gYR, dims = 1)
wgEU_baseline_mean = sum(wgE_baseline .* gYE, dims = 1) + sum(wgU_baseline .* gYU, dims = 1)
wgEU_stochastic_mean = sum(wgE_stochastic .* gYE, dims = 1) + sum(wgU_stochastic .* gYU, dims = 1)

τB = range(0, 20, step = 0.25)
wgB_baseline_mean = exp.((m.r + ϕ) .* (τB .- τB[end]))' .* wgEU_baseline_mean[1]
wgB_stochastic_mean = exp.((m.r + ϕ) .* (τB .- τB[end]))' .* wgEU_stochastic_mean[1]

p = plot(vcat(τB, τs_young, τs_old), [dp .* vcat(wgB_baseline_mean', wgEU_baseline_mean', wgR_baseline_mean') ./ Ybar  dp .* vcat(wgB_stochastic_mean', wgEU_stochastic_mean', wgR_stochastic_mean') ./ Ybar], linewidth = [2.5 2.5], linestyle = [:solid :dash], color = [:black :grey], lab = ["Discounting by interest rate"  "Discounting by individual MRS" ], xlabel = "Initial age", ylabel = "Welfare gain scaled by average labor income", xtickfontsize=12, ytickfontsize=12, guidefontsize = 12, legendfontsize = 12, legendfonthalign = :left, legend = :topleft, xticks = [0, 10, 20, 30, 40, 50, 60, 70, 80])
hline!([0], linestyle=:solid, color = :black, lab = "")
savefig(p, "$(MyPath)/Figures/FigureA8c.pdf") 


# compute the model using different calibration
println("4. Solving the model with alternative calibrations")
open("$(MyPath)/Tables/TableA8.tex", "w") do myfile
    wg_baseline, wg_stochastic, wg_rmse, wg_corr = solve_experiment()
    @printf(myfile, "Baseline  & %.2f & %.2f & %.2f  & %.2f \\\\ \n",dp * wg_baseline ./ Ybar,dp * wg_stochastic ./ Ybar,dp * wg_rmse ./ Ybar, wg_corr)
    experiments = [
    ("\$ \\sigma\$", m.σ, 0.2, :σ=> 0.2),
    ("\$ \\chi_U\$", m.χU, 0.3, :χU=> 0.3),
    ("\$ \\gamma \$", m.γ, 3.0, :γ=> 3.0),
    ("\$ \\phi \$", exp(-ϕ), exp(-0.01), :ϕ=> 0.01)
    ]
    for experiment in experiments
        println(experiment[4])
        local wg_baseline, wg_stochastic, wg_rmse, wg_corr = solve_experiment(; Dict(experiment[4])...)
        @printf(myfile, "%s: %.2f \$ \\rightarrow \$ %.2f & %.2f  & %.2f & %.2f & %.2f \\\\ \n", experiment[1], experiment[2], experiment[3], dp * wg_baseline / Ybar, dp * wg_stochastic / Ybar, dp * wg_rmse / Ybar, wg_corr)
    end
end

